//      gene_em.h
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

// CHANGE LOG
// 2012-04-02 : Commenting

#ifndef __GENE_EM_H
#define __GENE_EM_H

#include <boost/program_options.hpp>
#include <boost/foreach.hpp>
#include <map>
#include "stl-helper.h"

/** A hidden state template for use in EM algorithm
 *  it should implement a lessthan class to compare states 
 *  (for use in map)
 */
struct MotifHiddenState
{
  struct lessthan
  {
    bool operator() (const MotifHiddenState& a, const MotifHiddenState& b) const;
  };

  int rbs_inside; // 1 if there is RBS, 0 otherwise
  int rbs_position; // 0 if there is no RBS, otherwise it is the position

  std::string toString() const;
};

/** The (joint) probability map P(observed, hidden) */
template<typename H> 
class ProbMap : public std::map<H, double, typename H::lessthan>
{
 public:
  typedef typename std::map<H, double, typename H::lessthan> __BASE;
  typedef typename __BASE::value_type value_type;
 public:
  /** normalize the joint probability to posterior probability */
  void normalize();
  /** compute the marginal probability */
  double sum() const;
  /** print the probability map */
  void print() const;
};

/** Converting ACGT to/from 0123 */
class ACGT
{
	static std::string acgt;
 public:
	/** from char to int */
  static int c(char ch); 
  /** from int to char */ 
  static char ch(int c);
};

/** A (hidden) probabilistic model that could be trained with EM algorithm
 *  It should define prob_t (joint probability map) and datum_t (data type).
 *  Besides, it has to implement 
 *    prob_t compute_full_probability(const datum_t& x)
 *    void accumulate(const datum_t& x, const prob_t& p);
 * 	  void begin_accumulate();
 *    void end_accumulate();
 *  see below for more information
 */
class MotifModel
{
 public:
  struct Parameters {
    vd_t rbs_prob; // mixture probability of RBS model, (1-rbs_prob) = probability of no-RBS model
    vd_t rbs_position_prob;
    vd_t background_freq;
    vvd_t position_freq;

    std::string toString() const;
    void save(const std::string& outputFile) const;
  };

 public:
  typedef MotifHiddenState hidden_t;
  typedef ProbMap<hidden_t> prob_t; // needed by EM
  typedef std::string datum_t;      // needed by EM
  typedef Parameters param_t;

 public:
  /** Compute full (joint) probability map for a datum */
  prob_t compute_full_probability(const datum_t& x);

  /** The EM algorithm need accumulate(x,p), begin_accumulate() and end_accumulate() 
   *  print() is for debug
   */
  void accumulate(const datum_t& x, const prob_t& p);
  void begin_accumulate();
  void end_accumulate();
  std::string toString() const;

  /** random_initialize() and param() is used in main program to initialize and save model */
  void random_initialize(int motif_len, int seq_len);
  const param_t& param() const { return param_; }

 protected:
  prob_t create_prob_map(int len);
  double compute_background(const datum_t& x);
  double compute_motif(const datum_t& x, int pos, double background);
  int motif_len() const { return sz(param_.position_freq); }

 protected:
  param_t param_;
  param_t accumulate_temp_;

};


class PositionWeight
{
  vvd_t position_freq_;
  std::map<std::string, vd_t> background_markov_freq_;
  vd_t rbs_pos_freq_;
  
  double pos_score(const std::string& seq, int pos) const;
public:  
  void read_pwm(const std::string& pwmFile);
  double score(const std::string& seq) const;
  std::string toString() const;
  int motif_len() const { return sz(position_freq_); }
};

class Motif2ndModel
{
 public:
  struct Parameters {
    vd_t rbs_prob; // mixture probability of RBS model, (1-rbs_prob) = probability of no-RBS model
    vd_t rbs_position_prob;
    std::map<std::string, vd_t> background_markov_freq;
    vvd_t position_freq;

    std::string toString() const;
    void save(const std::string& outputFile) const;
  };

 public:
  typedef MotifHiddenState hidden_t;
  typedef ProbMap<hidden_t> prob_t; // needed by EM
  typedef std::string datum_t;      // needed by EM
  typedef Parameters param_t;

 public:
  /** Compute full (joint) probability map for a datum */
  prob_t compute_full_probability(const datum_t& x);

  /** The EM algorithm need accumulate(x,p), begin_accumulate() and end_accumulate() 
   *  print() is for debug
   */
  void accumulate(const datum_t& x, const prob_t& p);
  void begin_accumulate();
  void end_accumulate();
  std::string toString() const;

  /** random_initialize() and param() is used in main program to initialize and save model */
  void random_initialize(int motif_len, int seq_len);
  const param_t& param() const { return param_; }

 protected:
  prob_t create_prob_map(int len);
  double compute_background(const datum_t& x);
  double compute_motif(const datum_t& x, int pos, double background);
  int motif_len() const { return sz(param_.position_freq); }

 protected:
  param_t param_;
  param_t accumulate_temp_;

};

#endif
